import torch
from matplotlib import pyplot
from log import logger


class Train:

    def __init__(self, sequential, optimizer, scheduler, train_iter, epochs, device, accumulator):
        self.sequential = sequential
        self.optimizer = optimizer
        self.train_iter = train_iter
        self.epochs = epochs
        self.device = device
        self.accumulator = accumulator
        self.scheduler = scheduler

    def train(self):
        self.sequential.train(mode=True)
        self.sequential.to(self.device)
        accuracies = []
        losses = []
        for e in range(self.epochs):
            logger.info("----------------------")
            for x, y in self.train_iter:
                self.optimizer.zero_grad()
                x, y = x.to(self.device), y.to(self.device)
                y_pred = self.sequential(x)
                l = self.loss(x, y, y_pred)
                l.backward()
                self.optimizer.step()
                ev = self.train_accuracy(x, y, y_pred).item()
                print("\r", f"loss: {l.item()} ", f"lr: {self.scheduler.get_last_lr()} ", f"evaluation: {ev}", end='')
                logger.info(f"epochs: {e}, loss: {l.item()} , lr:{self.scheduler.get_last_lr()}")
                self.accumulator.collect(l.item(), ev)
            self.scheduler.step()
            accuracies.append(self.accumulator.ap())
            avg_loss = self.accumulator.al()
            print('     epoch ' + str(e) + ' avg loss: ' + str(avg_loss) + ' accuracy: ' + str(self.accumulator.ap()))
            logger.info(f"epochs {e}, avg loss: {avg_loss}, accuracy: {self.accumulator.ap()}")
            self.accumulator.reset()
            losses.append(avg_loss)

    def loss(self, x, y, y_pred):
        return torch.empty()

    def train_accuracy(self, x, y, y_pred):
        return 0.0


class Test:
    def __init__(self, sequential, test_iter, device, epochs, accumulator):
        self.sequential = sequential
        self.test_iter = test_iter
        self.accumulator = accumulator
        self.device = device
        self.epochs = epochs

    def test_accuracy(self, y, y_pred):
        return 0.0

    def test(self, mode='show', fn=None):
        self.sequential.eval()
        self.sequential.to(self.device)
        accuracies = []
        for e in range(self.epochs):
            for x, y in self.test_iter:
                x, y = x.to(self.device), y.to(self.device)
                y_pred = self.sequential(x)
                ev = self.test_accuracy(y, y_pred)
                self.accumulator.collect(e=ev)

            accuracies.append(self.accumulator.ap())
            self.accumulator.reset()

        test_fig = pyplot.figure()

        pyplot.title('dog_classify test')
        pyplot.plot(range(self.epochs), accuracies, 'b--', label='accuracy')
        pyplot.xlabel('epoch')
        pyplot.ylabel('accuracies')
        pyplot.legend()
        if mode == 'show':
            pyplot.show()
        elif mode == 'save':
            pyplot.savefig(fn)
